package com.kujie.ai.controller;

import com.kujie.ai.utils.TimeUtil;
import io.modelcontextprotocol.client.McpAsyncClient;
import io.modelcontextprotocol.server.McpServerFeatures;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.mcp.AsyncMcpToolCallbackProvider;
import org.springframework.ai.mcp.McpToolUtils;
import org.springframework.ai.mcp.SyncMcpToolCallbackProvider;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.http.MediaType;
import org.springframework.util.StreamUtils;
import org.springframework.web.bind.annotation.*;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.time.format.TextStyle;
import java.time.temporal.WeekFields;
import java.util.*;
import java.util.stream.Collectors;

@Slf4j
@RestController
@RequestMapping("/ai")
@CrossOrigin(origins = "*")
@RequiredArgsConstructor
public class ChatController {

    private final ChatModel chatModel;
    private final AsyncMcpToolCallbackProvider asyncMcpToolCallbackProvider;
    private final List<McpAsyncClient>  mcpAsyncClients;

    @GetMapping(value = "/steamChat", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public String steamChat(@RequestParam("input") String input) {
        // 1. 获取当前详细时间信息
        LocalDateTime now = LocalDateTime.now();
        String formattedTime = now.format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss"));
        String dayOfWeek = now.getDayOfWeek().getDisplayName(TextStyle.FULL, Locale.CHINESE);
        String season = TimeUtil.getCurrentSeason(now);
        int weekOfYear = now.get(WeekFields.of(Locale.CHINA).weekOfWeekBasedYear());


        Set<String> toolNames = Arrays.stream(asyncMcpToolCallbackProvider.getToolCallbacks()).map(toolCallback -> {
            McpServerFeatures.AsyncToolSpecification asyncToolSpecification = McpToolUtils.toAsyncToolSpecification(toolCallback);
            return asyncToolSpecification.tool().name();
        }).collect(Collectors.toSet());

        OpenAiChatOptions openAiChatOptions = OpenAiChatOptions.builder()
                .toolCallbacks(asyncMcpToolCallbackProvider.getToolCallbacks())
                .toolNames(toolNames)
                .toolChoice("auto")
                .toolNames(toolNames)
                .internalToolExecutionEnabled(true)
                .build();

        String systemPrompt = null;
        try {
            ClassPathResource resource = new ClassPathResource("template/system-prompt.txt");
            systemPrompt = StreamUtils.copyToString(resource.getInputStream(), StandardCharsets.UTF_8);
        } catch (IOException e) {
            log.error("提示词提取失败", e);
        }

        SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate(systemPrompt);
        String systemMessage = systemPromptTemplate.render(Map.of("time", formattedTime, "day", dayOfWeek, "season", season, "week", String.valueOf(weekOfYear), "city", "长沙"));

        Prompt prompt = new Prompt(
                List.of(
                        new SystemMessage(systemMessage),
                        new UserMessage(input)
                ),
                openAiChatOptions
        );

        ToolCallingManager toolCallingManager = ToolCallingManager.builder().build();

        ChatResponse chatResponse = chatModel.call(prompt);

        while (chatResponse.hasToolCalls()) {
            ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, chatResponse);
            prompt = new Prompt(toolExecutionResult.conversationHistory(), openAiChatOptions);

            chatResponse = chatModel.call(prompt);
        }

        return chatResponse.getResult().getOutput().getText();
    }
}
